use crate::chunk::Chunk;
use crate::page::Pages;

use alloc::alloc::Layout;
use alloc::vec::Vec;
use core::ptr::NonNull;

pub struct MemoryManager {
    mem: Vec<NonNull<u8>>,
    limit: Vec<NonNull<u8>>,
}

impl MemoryManager {
    // small chunks (size < 256) are divided into 32 bins, each bin has 8 bytes step
    const EXACT_BINS_CNT: usize = 32;
    // When a chunk is larger than 0x10000, it goes into the last bin
    // 256 <= size < 0x10000 are divided into 63 bins
    // Hierarchical Partitioning in get index.
    const SORTED_BINS_CNT: usize = 64;
    // 32 + 64 = 96
    pub const BINS_LEN: usize = Self::EXACT_BINS_CNT + Self::SORTED_BINS_CNT;
    // The minimum size of a chunk
    const MIN_DATA_SIZE: usize = Chunk::MIN_SIZE - Chunk::META_SIZE; // 24

    /// Get the index of the bin for the given size.
    pub fn get_index(size: usize) -> usize {
        let size = usize::max(size, Self::MIN_DATA_SIZE);
        // When size < 256, it goes into the exact bins
        if (size >> 8) == 0 {
            size / 8
        }
        // When size >= 0x10000, it goes into the last bin
        else if (size >> 8) >= 0x10000 {
            Self::BINS_LEN - 1
        } else {
            // m in [0, 16]
            // m is the highest valid bit of size/256
            let m = (usize::BITS - (size >> 8).leading_zeros() - 1) as usize;
            // each m has 4 bins
            // 32 + (m << 2) + (size >> (m + 6) & 3)
            Self::EXACT_BINS_CNT + (m << 2) + ((size >> (m + 6)) & 3)
        }
    }

    /// Create a new memory manager from the given memory.
    pub unsafe fn new(pages: &mut Pages) -> Self {
        /* let ptr_range = mem.as_mut_ptr_range();
        Self {
            mem: unsafe { NonNull::new_unchecked(ptr_range.start) },
            limit: unsafe { NonNull::new_unchecked(ptr_range.end) },
        } */
        let mut mem = Vec::new();
        let mut limit = Vec::new();
        for page in pages.0.iter_mut() {
            mem.push(NonNull::new_unchecked(page.as_mut_ptr()));
            limit.push(NonNull::new_unchecked(page.as_mut_ptr().add(page.size())));
        }
        Self { mem, limit }
    }

    /// Set a chunk in the bin.
    /// If the chunk is None, set the bin to null.
    /// Otherwise, set the bin to the data pointer of the chunk.
    unsafe fn set_bin_chunk(&mut self, index: usize, chunk: Option<Chunk>) {
        let ptr = if let Some(chunk) = chunk {
            chunk.data.as_ptr()
        } else {
            core::ptr::null_mut()
        };
        unsafe { *(self.mem[index].as_ptr().cast::<*mut u8>()) = ptr };
    }

    /// Get the start chunk.
    unsafe fn start_chunk(&self) -> Chunk {
        Chunk::new(
            NonNull::new_unchecked(self.mem[Self::BINS_LEN].as_ptr()),
            self.limit[Self::BINS_LEN],
        )
    }

    /// Get a chunk at the given index.
    /// If the chunk is null, return None.
    unsafe fn get_bin_chunk(&self, index: usize) -> Option<Chunk> {
        let chunk_addr = unsafe { *(self.mem[index].as_ptr().cast::<*mut u8>()) };
        if let Some(ptr) = NonNull::new(chunk_addr) {
            Some(Chunk::new(ptr, self.limit[index]))
        } else {
            None
        }
    }

    // add a chunk
    unsafe fn add_chunk(&mut self, mut chunk: Chunk) {
        let chunk_size = unsafe { chunk.get_size() };
        let index = Self::get_index(chunk_size - Chunk::META_SIZE);
        let mut bin_chunk = unsafe { self.get_bin_chunk(index) };
        if bin_chunk.is_none() {
            unsafe { self.set_bin_chunk(index, Some(chunk)) }
            for index in index + 1..Self::BINS_LEN {
                bin_chunk = unsafe { self.get_bin_chunk(index) };
                if bin_chunk.is_some() {
                    break;
                }
            }
        }
        let mut bin_chunk = bin_chunk.expect("top chunk always reachable from bins");
        while unsafe { !bin_chunk.is_top() && bin_chunk.get_size() <= chunk_size } {
            bin_chunk = if let Some(next_chunk) = unsafe { bin_chunk.get_next() } {
                next_chunk
            } else {
                break;
            };
        }
        unsafe {
            chunk.set_next(Some(bin_chunk));
            chunk.set_prev(bin_chunk.get_prev());
            bin_chunk.set_prev(Some(chunk));
            if let Some(mut prev_chunk) = chunk.get_prev() {
                prev_chunk.set_next(Some(chunk))
            }
        }
    }

    /// Initialize the memory manager with the given length.
    /// The length is the size of the memory of each page.
    pub unsafe fn init(&mut self, len: usize) {
        let s = 3; // Self::get_index(Self::MIN_DATA_SIZE);
        let e = Self::BINS_LEN;

        for index in s..e {
            unsafe { self.set_bin_chunk(index, None) }
        }

        unsafe {
            let mut chunk = self.start_chunk();
            let chunk_size = self.mem[Self::BINS_LEN]
                .as_ptr()
                .add(len)
                .offset_from(chunk.data.as_ptr()) as usize
                - Chunk::MIN_SIZE;
            chunk.set_cur_status(false, chunk_size);
            chunk.set_lower_in_use(true);

            let mut top_chunk = chunk.get_higher_chunk();
            top_chunk.set_cur_status(false, Chunk::MIN_SIZE);
            top_chunk.set_next(None);
            top_chunk.set_prev(None);
            self.set_bin_chunk(Self::get_index(usize::MAX), Some(top_chunk));

            self.add_chunk(chunk);
        }
    }

    // remove a chunk
    unsafe fn remove_chunk(&mut self, chunk: Chunk) {
        let mut next_chunk;

        next_chunk = chunk.get_next().unwrap();
        next_chunk.set_prev(chunk.get_prev());
        if let Some(mut prev_chunk) = chunk.get_prev() {
            prev_chunk.set_next(Some(next_chunk));
        }

        let index = Self::get_index(chunk.get_size() - Chunk::META_SIZE);
        if self.get_bin_chunk(index) == Some(chunk) {
            let chunk = if Self::get_index(next_chunk.get_size()) == index {
                Some(next_chunk)
            } else {
                None
            };
            self.set_bin_chunk(index, chunk);
        }
    }

    // `remove_chunk(top)` + `add_chunk(new_top)`
    unsafe fn update_top_chunk(&mut self, top: Chunk, mut new_top: Chunk) {
        new_top.set_next(None);
        new_top.set_prev(top.get_prev());
        if let Some(mut prev_chunk) = top.get_prev() {
            prev_chunk.set_next(Some(new_top));
        }
        let index = Self::get_index(usize::MAX);
        if self.get_bin_chunk(index) == Some(top) {
            self.set_bin_chunk(index, Some(new_top));
        }
    }

    unsafe fn find_smallest(&self, min_size: usize) -> Chunk {
        let mut chunk = None;
        for index in Self::get_index(min_size)..Self::BINS_LEN {
            chunk = unsafe { self.get_bin_chunk(index) };
            if chunk.is_some() {
                break;
            }
        }
        chunk.unwrap()
    }

    unsafe fn alloc(&mut self, layout: Layout) -> Result<NonNull<u8>, Chunk> {
        if layout.size() == 0 {
            return Ok(NonNull::dangling()); // feels like better than null?
        }
        let mut chunk = unsafe { self.find_smallest(layout.size()) };
        let mut user_data = unsafe { chunk.get_user_data(layout) };
        while user_data.is_none() {
            chunk = if let Some(next_chunk) = unsafe { chunk.get_next() } {
                next_chunk
            } else {
                return Err(chunk);
            };
            user_data = unsafe { chunk.get_user_data(layout) };
        }
        if unsafe { chunk.is_top() } {
            return Err(chunk); // top chunk is never used
        } else {
            unsafe {
                self.remove_chunk(chunk);
                if let Some(remain) = chunk.split(layout) {
                    self.add_chunk(remain)
                }
            }
        }
        unsafe { chunk.set_cur_status(true, chunk.get_size()) }
        let user_data = user_data.unwrap();
        let padding = unsafe { user_data.as_ptr().offset(-8) };
        let padding_size = unsafe { padding.offset_from(chunk.data.as_ptr()) } as usize;
        if padding_size != 0 {
            unsafe { *padding.cast::<u64>() = padding_size as u64 }
        }
        Ok(user_data)
    }

    unsafe fn dealloc(&mut self, user_data: *mut u8, layout: Layout) {
        let index = Self::get_index(layout.size());
        let mut chunk = unsafe { Chunk::from_user_data(user_data, layout, self.limit[index]) };
        if let Some(mut free_lower) = unsafe { chunk.get_free_lower_chunk() } {
            unsafe {
                self.remove_chunk(free_lower);
                chunk.set_cur_status(false, chunk.get_size());
                free_lower.coalesce(chunk);
                chunk = free_lower;
            }
        } else {
            unsafe { chunk.set_cur_status(false, chunk.get_size()) }
        }
        if let Some(free_higher) = unsafe { chunk.get_free_higher_chunk() } {
            if unsafe { !free_higher.is_top() } {
                unsafe {
                    self.remove_chunk(free_higher);
                    chunk.coalesce(free_higher);
                }
            }
        }
        unsafe { self.add_chunk(chunk) }
    }

    unsafe fn realloc(
        &mut self,
        user_data: *mut u8,
        layout: Layout,
        new_size: usize,
    ) -> Option<NonNull<u8>> {
        let index = Self::get_index(layout.size());
        let mut chunk = unsafe { Chunk::from_user_data(user_data, layout, self.limit[index]) };
        let new_layout = Layout::from_size_align(new_size, layout.align()).unwrap();
        if let Some(user_data) = unsafe { chunk.get_user_data(new_layout) } {
            return Some(user_data);
        }
        let Some(free_higher) = (unsafe { chunk.get_free_higher_chunk() }) else {
            return None;
        };
        if unsafe { free_higher.is_top() }
            || unsafe { chunk.get_size() + free_higher.get_size() } < new_size + Chunk::META_SIZE
        {
            return None;
        }
        unsafe {
            self.remove_chunk(free_higher);
            chunk.coalesce(free_higher);
        }
        if let Some(user_data) = unsafe { chunk.get_user_data(new_layout) } {
            let remain = unsafe { chunk.split(new_layout) };
            if let Some(remain) = remain {
                unsafe { self.add_chunk(remain) }
            }
            Some(user_data)
        } else {
            None
        }
    }

    pub unsafe fn alloc_in_mem(pages: &mut Pages, layout: Layout) -> *mut u8 {
        let mut mem_manger = Self::new(pages);
        let user_data = match unsafe { mem_manger.alloc(layout) } {
            Ok(user_data) => user_data.as_ptr(),
            Err(mut top) => {
                let top_index = Self::get_index(top.get_size() - Chunk::META_SIZE);
                let size = pages.size();
                match pages.grow(size + layout.size() + layout.align() + Chunk::META_SIZE) {
                    Ok(_) => {
                        mem_manger = Self::new(pages);
                        top.limit = mem_manger.limit[top_index];
                        let new_size = pages.size();
                        unsafe {
                            let mut new_top = Chunk::new(
                                NonNull::new_unchecked(
                                    pages[top_index].as_mut_ptr_range().end.sub(Chunk::MIN_SIZE),
                                ),
                                mem_manger.limit[Self::BINS_LEN],
                            );
                            new_top.set_prev(None);
                            new_top.set_next(None);
                            new_top.set_cur_status(false, Chunk::MIN_SIZE);
                            mem_manger.update_top_chunk(top, new_top);
                            top.set_cur_status(false, new_size - size);
                            if let Some(mut free_lower) = top.get_free_lower_chunk() {
                                mem_manger.remove_chunk(free_lower);
                                free_lower
                                    .set_cur_status(false, free_lower.get_size() + top.get_size());
                                mem_manger.add_chunk(free_lower);
                            } else {
                                mem_manger.add_chunk(top);
                            }
                            mem_manger.alloc(layout)
                        }
                        .unwrap()
                        .as_ptr()
                    }
                    Err(_) => core::ptr::null_mut(),
                }
            }
        };
        user_data
    }

    pub unsafe fn dealloc_in_space(pages: &mut Pages, user_data: *mut u8, layout: Layout) {
        let mut mem_manger = Self::new(pages);
        unsafe {
            mem_manger.dealloc(user_data, layout);
        }
    }

    pub unsafe fn realloc_in_space(
        pages: &mut Pages,
        user_data: *mut u8,
        layout: Layout,
        new_size: usize,
    ) -> *mut u8 {
        let mut mem_manger = Self::new(pages);
        if let Some(user_data) = unsafe { mem_manger.realloc(user_data, layout, new_size) } {
            return user_data.as_ptr();
        }
        let new_user_data = unsafe {
            Self::alloc_in_mem(
                pages,
                Layout::from_size_align(new_size, layout.align()).unwrap(),
            )
        };
        if new_user_data.is_null() {
            core::ptr::null_mut()
        } else {
            unsafe {
                core::ptr::copy_nonoverlapping(user_data, new_user_data, layout.size());
                Self::dealloc_in_space(pages, user_data, layout);
            }
            new_user_data
        }
    }
}
